import os
import uuid
import shutil
from datetime import datetime
from werkzeug.utils import secure_filename
import fitz
from docx import Document as DocxDocument
from langdetect import detect
from paddleocr import PaddleOCR
from pymongo import MongoClient
from config.config import Config
import re
import logging

logger = logging.getLogger(__name__)

class DocumentUploader:
    def __init__(self):
        self.config = Config()
        self.client = MongoClient(self.config.MONGODB_URI)
        self.db = self.client[self.config.DB_NAME]
        try:
            self.ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)
        except Exception as e:
            logger.warning(f"OCR initialization failed: {e}. OCR features will be disabled.")
            self.ocr = None
        
    def allowed_file(self, filename):
        return '.' in filename and \
               filename.rsplit('.', 1)[1].lower() in self.config.ALLOWED_EXTENSIONS
    
    def detect_language(self, text_sample):
        try:
            if not text_sample or len(text_sample.strip()) < 10:
                return 'english'
            detected = detect(text_sample[:1000])
            language_map = {
                'en': 'english',
                'zh-cn': 'chinese',
                'zh-tw': 'chinese',
                'pt': 'portuguese',
                'ja': 'japanese'
            }
            return language_map.get(detected, 'english')
        except Exception as e:
            logger.warning(f"Language detection failed: {e}")
            return 'english'
    
    def extract_pdf_content(self, file_path):
        """Extract text and images from PDF"""
        logger.debug(f"Extracting PDF content from {file_path}")
        doc = fitz.open(file_path)
        content = {
            'text': '',
            'images': [],
            'page_count': len(doc)
        }
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()
            content['text'] += f"\n--- Page {page_num + 1} ---\n{text}"
            
            # Skip OCR if not available
            if not self.ocr:
                continue
                
            image_list = page.get_images()
            for img_index, img in enumerate(image_list[:5]):  # Limit to 5 images per page
                try:
                    xref = img[0]
                    base_image = doc.extract_image(xref)
                    image_bytes = base_image["image"]
                    
                    img_path = f"/tmp/temp_img_{uuid.uuid4().hex}.png"
                    with open(img_path, "wb") as img_file:
                        img_file.write(image_bytes)
                    
                    ocr_result = self.ocr.ocr(img_path, cls=True)
                    ocr_text = ""
                    if ocr_result and ocr_result[0]:
                        for line in ocr_result[0]:
                            ocr_text += line[1][0] + "\n"
                    
                    content['images'].append({
                        'page_number': page_num + 1,
                        'image_index': img_index,
                        'ocr_text': ocr_text
                    })
                    
                    if os.path.exists(img_path):
                        os.remove(img_path)
                except Exception as e:
                    logger.warning(f"Error processing image on page {page_num + 1}: {e}")
        
        doc.close()
        logger.debug(f"PDF extraction complete: {len(content['text'])} characters extracted")
        return content
    
    def extract_docx_content(self, file_path):
        """Extract text and images from DOCX"""
        logger.debug(f"Extracting DOCX content from {file_path}")
        doc = DocxDocument(file_path)
        content = {
            'text': '',
            'images': [],
            'page_count': len(doc.paragraphs)
        }
        
        for para in doc.paragraphs:
            content['text'] += para.text + "\n"
        
        logger.debug(f"DOCX extraction complete: {len(content['text'])} characters extracted")
        return content
    
    def extract_txt_content(self, file_path):
        """Extract text from TXT file"""
        logger.debug(f"Extracting TXT content from {file_path}")
        content = {
            'text': '',
            'images': [],
            'page_count': 1
        }
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content['text'] = f.read()
        except UnicodeDecodeError:
            # Try with different encoding
            try:
                with open(file_path, 'r', encoding='latin-1') as f:
                    content['text'] = f.read()
            except Exception as e:
                logger.error(f"Failed to read text file: {e}")
                content['text'] = 'Could not read text content'
        
        logger.debug(f"TXT extraction complete: {len(content['text'])} characters extracted")
        return content

    def _sanitize_filename(self, filename):
        """Sanitize filename to only contain allowed characters"""
        if not filename:
            return ""
        sanitized = re.sub(r'[^A-Za-z0-9._-]', '', filename)
        if len(sanitized) > 128:
            sanitized = sanitized[:128]
        return sanitized

    def _parse_date(self, date_str):
        """Parse date string to datetime object"""
        if not date_str:
            return None
        try:
            # Try yyyy/mm/dd format
            if '/' in date_str:
                return datetime.strptime(date_str, '%Y/%m/%d')
            # Try yyyy-mm-dd format
            elif '-' in date_str:
                return datetime.strptime(date_str, '%Y-%m-%d')
            else:
                logger.warning(f"Unexpected date format: {date_str}")
                return None
        except ValueError as e:
            logger.error(f"Date parsing error for '{date_str}': {e}")
            return None

    def _validate_and_process_tags(self, tags_input):
        """Validate and process tags"""
        tags = []
        if not tags_input:
            return tags
            
        if isinstance(tags_input, str):
            tag_list = [tag.strip() for tag in tags_input.split(',') if tag.strip()]
        elif isinstance(tags_input, list):
            tag_list = tags_input
        else:
            return tags
        
        for tag in tag_list:
            if len(tag) > 32:
                logger.warning(f"Tag too long: {tag}")
                continue
                
            if not re.match(r'^[A-Za-z0-9]+$', tag):
                logger.warning(f"Invalid tag format: {tag}")
                continue
                
            tags.append(tag)
            self._add_tag_to_list(tag)
        
        return tags[:8]

    def _add_tag_to_list(self, tag):
        """Add tag to tags collection"""
        try:
            existing_tag = self.db.tags.find_one({'name': tag})
            if not existing_tag:
                self.db.tags.insert_one({
                    'name': tag,
                    'usage_count': 1,
                    'created_at': datetime.now()
                })
            else:
                self.db.tags.update_one(
                    {'name': tag},
                    {'$inc': {'usage_count': 1}}
                )
        except Exception as e:
            logger.error(f"Error adding tag to list: {e}")
    
    def _update_collection_count(self, collection_name):
        """Update document count for collection"""
        try:
            document_count = self.db.documents.count_documents({'collection': collection_name})
            self.db.collections.update_one(
                {'name': collection_name},
                {'$set': {'document_count': document_count}},
                upsert=True
            )
            logger.debug(f"Updated collection '{collection_name}' count to {document_count}")
        except Exception as e:
            logger.error(f"Error updating collection count: {e}")
    
    def upload_document(self, file, metadata):
        """Main upload function"""
        try:
            logger.info(f"Starting document upload: {file.filename}")
            logger.debug(f"Metadata: {metadata}")
            
            # Validate file type
            if not self.allowed_file(file.filename):
                logger.error(f"File type not allowed: {file.filename}")
                return {'status': 'error', 'message': 'File type not allowed'}
            
            # Get file extension
            original_extension = file.filename.rsplit('.', 1)[1].lower()
            logger.debug(f"File extension: {original_extension}")
            
            # Determine filename
            user_file_name = metadata.get('file_name', '').strip()
            if user_file_name:
                sanitized_name = self._sanitize_filename(user_file_name)
                logger.debug(f"Using user-provided filename: {sanitized_name}")
            else:
                sanitized_name = self._sanitize_filename(
                    file.filename.rsplit('.', 1)[0]
                )
                logger.debug(f"Using original filename: {sanitized_name}")
            
            if not sanitized_name:
                sanitized_name = str(uuid.uuid4().hex)[:12]
                logger.warning(f"Generated random filename: {sanitized_name}")
            
            # Add extension
            file_name_with_ext = f"{sanitized_name}.{original_extension}"
            
            # Ensure unique filename
            upload_path = os.path.join(self.config.UPLOAD_FOLDER, file_name_with_ext)
            counter = 1
            while os.path.exists(upload_path):
                file_name_with_ext = f"{sanitized_name}_{counter}.{original_extension}"
                upload_path = os.path.join(self.config.UPLOAD_FOLDER, file_name_with_ext)
                counter += 1
            
            logger.debug(f"Final upload path: {upload_path}")
            
            # Save file
            file.save(upload_path)
            logger.info(f"File saved to: {upload_path}")
            
            # Copy PDF to static folder for web access
            if original_extension == 'pdf':
                pdf_path = os.path.join(self.config.PDF_FOLDER, file_name_with_ext)
                shutil.copy(upload_path, pdf_path)
                logger.debug(f"PDF copied to static folder: {pdf_path}")
            
            # Extract content based on file type
            if original_extension == 'pdf':
                content = self.extract_pdf_content(upload_path)
            elif original_extension in ['docx', 'doc']:
                content = self.extract_docx_content(upload_path)
            elif original_extension == 'txt':
                content = self.extract_txt_content(upload_path)
            elif original_extension == 'xlsx':
                # Basic support for XLSX
                content = {
                    'text': 'Excel file - content extraction not fully supported',
                    'images': [],
                    'page_count': 1
                }
            else:
                content = {
                    'text': 'Content extraction not supported for this file type',
                    'images': [],
                    'page_count': 1
                }
            
            # Detect language
            detected_language = self.detect_language(content['text'])
            logger.debug(f"Detected language: {detected_language}")
            
            # Parse dates
            effective_date = self._parse_date(metadata.get('effective_date'))
            if not effective_date:
                logger.error("Failed to parse effective date")
                return {'status': 'error', 'message': 'Invalid effective date format'}
            
            # Process tags
            tags = self._validate_and_process_tags(metadata.get('tags'))
            logger.debug(f"Processed tags: {tags}")
            
            # Prepare document data
            current_time = datetime.now()
            document_data = {
                'filename': file_name_with_ext,
                'file_name': sanitized_name,
                'original_name': file.filename,
                'file_path': upload_path,
                'file_size': os.path.getsize(upload_path),
                'file_format': original_extension,
                'upload_time': current_time,
                'update_time': current_time,
                'country': metadata.get('country', ''),
                'document_type': metadata.get('document_type', ''),
                'regulation_id': metadata.get('regulation_id', ''),
                'issuing_authority': metadata.get('issuing_authority', ''),
                'effective_date': effective_date,
                'version': metadata.get('version', ''),
                'language': detected_language,
                'content_text': content['text'],
                'extracted_images': content['images'],
                'page_count': content['page_count'],
                'processing_status': 'completed',
                'last_updated': current_time,
                'collection': metadata.get('collection', ''),
                'tags': tags,
                'status': metadata.get('status', 'Draft')
            }
            
            logger.debug("Inserting document into database...")
            result = self.db.documents.insert_one(document_data)
            doc_id = str(result.inserted_id)
            logger.info(f"Document inserted with ID: {doc_id}")
            
            # Update collection count
            if metadata.get('collection'):
                self._update_collection_count(metadata.get('collection'))
            
            return {
                'status': 'success',
                'doc_id': doc_id,
                'filename': file_name_with_ext,
                'file_name': sanitized_name,
                'language': detected_language,
                'collection': metadata.get('collection'),
                'status_field': metadata.get('status', 'Draft'),
                'message': 'Document uploaded successfully'
            }
            
        except Exception as e:
            logger.error(f"Upload error: {str(e)}", exc_info=True)
            return {'status': 'error', 'message': f'Upload failed: {str(e)}'}